import logging
from pymilvus import connections, Collection, utility
from typing import List, Dict, Optional
import numpy as np

from rag_test.legal_embeddings import text_embeddings

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def hybrid_search_labor_law(
        query_text: str,
        collection_name: str = "labor_contract_law",
        filter_conditions: Optional[Dict] = None,
        top_k: int = 5,
        nprobe: int = 20,
        output_fields: Optional[List[str]] = None,
        host: str = "localhost",
        port: str = "19530"
) -> List[Dict]:
    """
    劳动合同法混合检索函数（完全兼容pymilvus 2.6.0）

    参数:
        query_text: 查询文本
        collection_name: 集合名称
        filter_conditions: 过滤条件，支持格式:
            - 简单匹配: {"field": value} (默认==匹配)
            - 模糊匹配: {"field": "*value*"} (自动转为like匹配)
            - 自定义操作: {"field": {"op": "like", "value": "xxx"}}
        top_k: 返回结果数量
        nprobe: 检索参数
        output_fields: 返回字段列表
        host: Milvus服务地址
        port: Milvus服务端口
    """
    # 1. 输入验证
    if not query_text or not query_text.strip():
        raise ValueError("查询文本不能为空")

    # 2. 默认返回字段
    if output_fields is None:
        output_fields = [
            "id", "law_name", "chapter", "chapter_code",
            "article_num", "article_num_code", "article_title",
            "content_clean", "applicable_scenarios", "retrieval_weight"
        ]
    # 确保id字段在返回字段中
    if "id" not in output_fields:
        output_fields.append("id")

    # 3. 生成查询向量
    try:
        query_input = [{"vector_text": query_text}]
        query_embedding = text_embeddings(query_input)[0]
        logger.info("查询向量生成成功")
    except Exception as e:
        raise RuntimeError(f"生成查询向量失败：{str(e)}")

    # 4. 连接Milvus并加载集合
    try:
        # 连接Milvus
        connections.connect(alias="default", host=host, port=port)
        logger.info(f"成功连接到Milvus服务 {host}:{port}")

        # 验证集合是否存在
        if not utility.has_collection(collection_name):
            raise ValueError(f"集合{collection_name}不存在，请先写入数据")

        # 加载集合
        collection = Collection(name=collection_name)
        collection.load()
        logger.info(f"已加载集合{collection_name}到内存")
    except Exception as e:
        raise ConnectionError(f"Milvus连接或集合加载失败：{str(e)}")

    # 5. 构建元数据过滤表达式
    filter_expr = ""
    if filter_conditions:
        conditions = []
        for field, value in filter_conditions.items():
            # 解析操作符和值
            if isinstance(value, dict) and "op" in value and "value" in value:
                op = value["op"]
                val = value["value"]
            else:
                # 自动判断匹配类型：含*则视为模糊匹配
                if isinstance(value, str) and '*' in value:
                    op = "like"
                    val = value.replace('*', '%')  # 转换为Milvus的like格式
                else:
                    op = "=="
                    val = value

            # 处理值的引号包裹
            if isinstance(val, str):
                # 转义单引号
                val = val.replace("'", "\\'")
                conditions.append(f"{field} {op} '{val}'")
            else:
                conditions.append(f"{field} {op} {val}")

        filter_expr = " and ".join(conditions)
        logger.info(f"元数据过滤条件：{filter_expr}")

    # 6. 执行混合检索
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": nprobe}
    }

    try:
        results = collection.search(
            data=[query_embedding],
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            expr=filter_expr,
            output_fields=output_fields
        )
        logger.info(f"检索完成，返回{len(results[0])}条结果")
    except Exception as e:
        raise RuntimeError(f"检索执行失败：{str(e)}")
    finally:
        connections.disconnect("default")
        logger.info("已断开Milvus连接")

    # 7. 解析检索结果
    parsed_results = []
    for hits in results:
        for hit in hits:
            # 计算相似度（L2距离转相似度）
            similarity = 1 - hit.distance if 0 <= hit.distance <= 1 else 0
            # 提取元数据
            metadata = {}
            for field in output_fields:
                if field == "id":
                    continue
                metadata[field] = hit.entity.get(field) or "N/A"  # 处理空值

            result_item = {
                "similarity": round(similarity, 4),
                "rank": len(parsed_results) + 1,
                "id": hit.entity.get("id"),
                "metadata": metadata
            }
            parsed_results.append(result_item)

    return parsed_results


# 测试函数
if __name__ == "__main__":
    # 测试1：基础语义检索（无过滤）
    print("=== 测试1：查询'未签劳动合同的补偿' ===")
    try:
        result1 = hybrid_search_labor_law(
            query_text="未签劳动合同有什么补偿",
            top_k=3
        )
        for item in result1:
            print(f"排名{item['rank']}（相似度{item['similarity']}）：")
            print(f"{item['metadata']['article_num']} {item['metadata']['article_title']}")
            print(f"内容：{item['metadata']['content_clean'][:100]}...\n")
    except Exception as e:
        logger.error(f"测试1失败：{e}")

    # 测试2：带元数据过滤的检索（使用模糊匹配）
    print("=== 测试2：限定第四章的经济补偿相关法条 ===")
    try:
        result2 = hybrid_search_labor_law(
            query_text="解除劳动合同的经济补偿",
            filter_conditions={
                "chapter_code": "04",
                "applicable_scenarios": "*compensation*"  # 使用*表示模糊匹配
            },
            top_k=2
        )
        for item in result2:
            print(f"排名{item['rank']}（相似度{item['similarity']}）：")
            print(f"{item['metadata']['article_num']} {item['metadata']['article_title']}")
            print(f"章节：{item['metadata']['chapter']}\n")
    except Exception as e:
        logger.error(f"测试2失败：{e}")

    # 测试3：自定义操作符过滤
    print("=== 测试3：使用自定义操作符 ===")
    try:
        result3 = hybrid_search_labor_law(
            query_text="工资标准",
            filter_conditions={
                "article_num": {"op": ">=", "value": "10"}  # 自定义>=操作
            },
            top_k=2
        )
        for item in result3:
            print(f"排名{item['rank']}（相似度{item['similarity']}）：")
            print(f"{item['metadata']['article_num']} {item['metadata']['article_title']}\n")
    except Exception as e:
        logger.error(f"测试3失败：{e}")